package ynu.edu.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import ynu.edu.domain.dto.Txt2ImgRequest;
import ynu.edu.domain.dto.Text2ImgParameters;
import ynu.edu.entity.ModelConfig;
import ynu.edu.service.ModelConfigService;

import javax.annotation.Resource;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

/**
 * 参数构造工具类 - 重构版，支持从ModelConfig查询所有参数
 */
@Slf4j
@Component
public class Text2imgConstructionUtil {

    @Resource
    private ModelConfigService modelConfigService;

    @Resource
    private ObjectMapper objectMapper;

    // 默认参数值
    private static final Integer DEFAULT_STEPS = 20;
    private static final String DEFAULT_SAMPLER_NAME = "DPM++ 2M";
    private static final BigDecimal DEFAULT_CFG_SCALE = new BigDecimal("7.0");
    private static final Integer DEFAULT_WIDTH = 512;
    private static final Integer DEFAULT_HEIGHT = 512;
    private static final Integer DEFAULT_BATCH_SIZE = 1;
    private static final String DEFAULT_SCHEDULER = "Karras";
    private static final BigDecimal DEFAULT_DENOISING_STRENGTH = null;
    private static final Boolean DEFAULT_RESTORE_FACES = false;
    private static final Boolean DEFAULT_TILING = false;
    private static final Boolean DEFAULT_ENABLE_HR = false;
    private static final BigDecimal DEFAULT_HR_SCALE = new BigDecimal("2.0");
    private static final String DEFAULT_HR_UPSCALER = null;
    private static final Integer DEFAULT_HR_SECOND_PASS_STEPS = 0;

    /**
     * 构造参数（基于Txt2ImgRequest）
     * @param request 文本生成图像请求
     * @return 构造的参数对象
     */
    public Text2ImgParameters constructParameters(Txt2ImgRequest request) {
        try {
            log.info("开始构造Text2Img参数，请求: {}", request);
            
            // 1. 获取模型配置
            ModelConfig modelConfig = modelConfigService.getModelForParams(request.getModelName());
            if (modelConfig == null) {
                log.warn("无法获取模型配置");
                return null;
            }
            
            log.info("获取到模型配置: {}", modelConfig.getModelName());
            
            // 2. 使用模型配置构造参数
            return constructParametersWithModelConfig(request, modelConfig);
            
        } catch (Exception e) {
            log.error("Text2Img参数构造失败", e);
            throw new RuntimeException("Text2Img参数构造失败: " + e.getMessage());
        }
    }

    /**
     * 使用模型配置构造参数
     */
    private Text2ImgParameters constructParametersWithModelConfig(Txt2ImgRequest request, ModelConfig modelConfig) {
        // 生成10位随机种子（如果模型配置中没有指定种子）
        Long seed = modelConfig.getSeed() != null ? modelConfig.getSeed() : generateRandomSeed();
        
        // 构造完整提示词
        String fullPrompt = constructFullPrompt(modelConfig.getPositivePrompt(), request.getPrompt());
        String fullNegativePrompt = constructFullNegativePrompt(modelConfig.getNegativePrompt(), request.getNegativePrompt());
        
        // 构建模型覆盖设置
        Map<String, String> overrideSettings = buildOverrideSettings(modelConfig);
        
        // 构建结果，优先使用模型配置中的值，其次使用请求中的值，最后使用默认值
        Text2ImgParameters result = new Text2ImgParameters();
        
        // 提示词相关
        result.setPrompt(fullPrompt);
        result.setNegativePrompt(fullNegativePrompt);
        result.setSeed(seed);
        
        // 尺寸相关 - 优先使用请求中的值，其次使用模型配置中的默认值
        result.setWidth(getValueOrDefault(request.getWidth(), modelConfig.getDefaultWidth(), DEFAULT_WIDTH));
        result.setHeight(getValueOrDefault(request.getHeight(), modelConfig.getDefaultHeight(), DEFAULT_HEIGHT));
        
        // 生成参数 - 优先使用模型配置中的值
        result.setSteps(getValueOrDefault(null, modelConfig.getSteps(), DEFAULT_STEPS));
        result.setSamplerName(getValueOrDefault(null, modelConfig.getSamplerName(), DEFAULT_SAMPLER_NAME));
        result.setScheduler(getValueOrDefault(null, modelConfig.getScheduler(), DEFAULT_SCHEDULER));
        result.setCfgScale(getValueOrDefault(null, modelConfig.getCfgScale(), DEFAULT_CFG_SCALE));
        result.setBatchSize(getValueOrDefault(null, modelConfig.getBatchSize(), DEFAULT_BATCH_SIZE));
        result.setDenoisingStrength(getValueOrDefault(null, modelConfig.getDenoisingStrength(), DEFAULT_DENOISING_STRENGTH));
        
        // 高级参数 - 使用模型配置中的值
        result.setRestoreFaces(getValueOrDefault(null, modelConfig.getRestoreFaces(), DEFAULT_RESTORE_FACES));
        result.setTiling(getValueOrDefault(null, modelConfig.getTiling(), DEFAULT_TILING));
        result.setEnableHr(getValueOrDefault(null, modelConfig.getEnableHr(), DEFAULT_ENABLE_HR));
        result.setHrScale(getValueOrDefault(null, modelConfig.getHrScale(), DEFAULT_HR_SCALE));
        result.setHrUpscaler(getValueOrDefault(null, modelConfig.getHrUpscaler(), DEFAULT_HR_UPSCALER));
        result.setHrSecondPassSteps(getValueOrDefault(null, modelConfig.getHrSecondPassSteps(), DEFAULT_HR_SECOND_PASS_STEPS));
        
        // 模型覆盖设置
        result.setOverrideSettings(overrideSettings);
        
        log.info("Text2Img参数构造完成，使用模型配置: {}, 种子: {}", modelConfig.getModelName(), seed);
        log.debug("构造的参数详情: {}", result);
        
        return result;
    }


    /**
     * 获取值，优先级：请求值 > 模型配置值 > 默认值
     */
    private <T> T getValueOrDefault(T requestValue, T modelConfigValue, T defaultValue) {
        if (requestValue != null) {
            return requestValue;
        }
        if (modelConfigValue != null) {
            return modelConfigValue;
        }
        return defaultValue;
    }

    /**
     * 生成10位随机种子
     */
    private Long generateRandomSeed() {
        Random random = new Random();
        // 生成10位随机数
        long seed = random.nextInt(900000000) + 100000000; // 100000000-999999999
        return seed;
    }

    /**
     * 构造完整正面提示词
     */
    private String constructFullPrompt(String modelPrompt, String userPrompt) {
        if (modelPrompt == null || modelPrompt.trim().isEmpty()) {
            return userPrompt;
        }
        
        if (userPrompt == null || userPrompt.trim().isEmpty()) {
            return modelPrompt.trim();
        }
        
        // 使用逗号分隔拼接
        return modelPrompt.trim() + ", " + userPrompt.trim();
    }

    /**
     * 构造完整负面提示词
     */
    private String constructFullNegativePrompt(String modelNegativePrompt, String userNegativePrompt) {
        if (modelNegativePrompt == null && userNegativePrompt == null) {
            return null;
        }
        
        if (modelNegativePrompt == null) {
            return userNegativePrompt;
        }
        
        if (userNegativePrompt == null) {
            return modelNegativePrompt;
        }
        
        // 使用逗号分隔拼接
        return modelNegativePrompt.trim() + ", " + userNegativePrompt.trim();
    }

    /**
     * 构建模型覆盖设置
     */
    private Map<String, String> buildOverrideSettings(ModelConfig modelConfig) {
        Map<String, String> overrideSettings = new HashMap<>();
        
        // 添加模型配置中的覆盖设置
        if (modelConfig != null && modelConfig.getModelFile() != null && !modelConfig.getModelFile().trim().isEmpty()) {
            overrideSettings.put("sd_model_checkpoint", modelConfig.getModelFile());
            log.debug("添加模型覆盖设置: sd_model_checkpoint = {}", modelConfig.getModelFile());
        }
        
        return overrideSettings;
    }

    /**
     * 获取默认参数值（用于调试和测试）
     */
    public Map<String, Object> getDefaultParameters() {
        Map<String, Object> defaults = new HashMap<>();
        defaults.put("steps", DEFAULT_STEPS);
        defaults.put("sampler_name", DEFAULT_SAMPLER_NAME);
        defaults.put("cfg_scale", DEFAULT_CFG_SCALE);
        defaults.put("width", DEFAULT_WIDTH);
        defaults.put("height", DEFAULT_HEIGHT);
        defaults.put("batch_size", DEFAULT_BATCH_SIZE);
        defaults.put("scheduler", DEFAULT_SCHEDULER);
        defaults.put("denoising_strength", DEFAULT_DENOISING_STRENGTH);
        defaults.put("restore_faces", DEFAULT_RESTORE_FACES);
        defaults.put("tiling", DEFAULT_TILING);
        defaults.put("enable_hr", DEFAULT_ENABLE_HR);
        defaults.put("hr_scale", DEFAULT_HR_SCALE);
        defaults.put("hr_upscaler", DEFAULT_HR_UPSCALER);
        defaults.put("hr_second_pass_steps", DEFAULT_HR_SECOND_PASS_STEPS);
        return defaults;
    }

    /**
     * 构建Txt2Img请求体
     * @param params 构造的参数对象
     * @return 请求体Map
     */
    public Map<String, Object> buildRequestBody(Text2ImgParameters params) {
        try {
            log.debug("开始构建Text2Img请求体，参数: {}", params);
            
            Map<String, Object> requestBody = new HashMap<>();
            
            // 必需参数
            requestBody.put("prompt", params.getPrompt() != null ? params.getPrompt() : "");
            requestBody.put("negative_prompt", params.getNegativePrompt() != null ? params.getNegativePrompt() : "");
            requestBody.put("steps", params.getSteps() != null ? params.getSteps() : DEFAULT_STEPS);
            requestBody.put("sampler_name", params.getSamplerName() != null ? params.getSamplerName() : DEFAULT_SAMPLER_NAME);
            requestBody.put("cfg_scale", params.getCfgScale() != null ? params.getCfgScale() : DEFAULT_CFG_SCALE);
            requestBody.put("seed", params.getSeed() != null ? params.getSeed() : -1);
            requestBody.put("width", params.getWidth() != null ? params.getWidth() : DEFAULT_WIDTH);
            requestBody.put("height", params.getHeight() != null ? params.getHeight() : DEFAULT_HEIGHT);
            requestBody.put("batch_size", params.getBatchSize() != null ? params.getBatchSize() : DEFAULT_BATCH_SIZE);
            
            // 可选参数
            if (params.getDenoisingStrength() != null) {
                requestBody.put("denoising_strength", params.getDenoisingStrength());
            }
            
            if (params.getScheduler() != null) {
                requestBody.put("scheduler", params.getScheduler());
            }
            
            // 高级参数
            if (params.getRestoreFaces() != null) {
                requestBody.put("restore_faces", params.getRestoreFaces());
            }
            
            if (params.getTiling() != null) {
                requestBody.put("tiling", params.getTiling());
            }
            
            if (params.getEnableHr() != null) {
                requestBody.put("enable_hr", params.getEnableHr());
            }
            
            if (params.getHrScale() != null) {
                requestBody.put("hr_scale", params.getHrScale());
            }
            
            if (params.getHrUpscaler() != null) {
                requestBody.put("hr_upscaler", params.getHrUpscaler());
            }
            
            if (params.getHrSecondPassSteps() != null) {
                requestBody.put("hr_second_pass_steps", params.getHrSecondPassSteps());
            }
            
            // 模型覆盖设置
            if (params.getOverrideSettings() != null && !params.getOverrideSettings().isEmpty()) {
                requestBody.put("override_settings", params.getOverrideSettings());
            }
            
            // 默认参数
            requestBody.put("send_images", true);
            requestBody.put("save_images", false);
            requestBody.put("do_not_save_samples", false);
            requestBody.put("do_not_save_grid", false);
            requestBody.put("override_settings_restore_afterwards", true);
            requestBody.put("disable_extra_networks", false);
            requestBody.put("firstphase_width", 0);
            requestBody.put("firstphase_height", 0);
            requestBody.put("hr_resize_x", 0);
            requestBody.put("hr_resize_y", 0);
            requestBody.put("hr_prompt", "");
            requestBody.put("hr_negative_prompt", "");
            requestBody.put("sampler_index", "Euler");
            requestBody.put("script_name", null);
            requestBody.put("script_args", new ArrayList<>());
            requestBody.put("alwayson_scripts", new HashMap<>());
            
            log.debug("Text2Img请求体构建完成，参数数量: {}", requestBody.size());
            
            // 输出发送到模型的完整参数JSON
            outputRequestBodyJson(requestBody);
            
            return requestBody;
            
        } catch (Exception e) {
            log.error("构建Text2Img请求体失败", e);
            throw new RuntimeException("构建Text2Img请求体失败: " + e.getMessage());
        }
    }

    /**
     * 输出请求体JSON
     */
    private void outputRequestBodyJson(Map<String, Object> requestBody) {
        try {
            // 转换为JSON并输出
            String jsonOutput = objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(requestBody);
            log.info("发送到Stable Diffusion模型的完整参数JSON:");
            log.info("{}", jsonOutput);
            
        } catch (Exception e) {
            log.warn("输出请求体JSON失败", e);
        }
    }
} 